import argparse
import pandas as pd
import torch
from origin_tcn import TCN as O_TCN
from utils import sigmoid
from data_load import data_eval
from index_factors import end_date

cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if cuda else 'cpu')

parser = argparse.ArgumentParser(description='Sequence Modeling - (Permuted) Sequential MNIST')
parser.add_argument('--batch_size', type=int, default=32, metavar='N',
                    help='batch size (default: 64)')
parser.add_argument('--cuda', action='store_false', default=True,
                    help='use CUDA (default: False)')
parser.add_argument('--dropout', type=float, default=0.2,
                    help='dropout applied to layers (default: 0.05)')
parser.add_argument('--ksize', type=int, default=7,
                    help='kernel size (default: 7)')
parser.add_argument('--levels', type=int, default=4,
                    help='# of levels (default: 8)')
parser.add_argument('--nhid', type=int, default=25,
                    help='number of hidden units per layer (default: 25)')
args = parser.parse_args()

n_classes = 1
input_channels = 13
seq_length = 300
steps = 0

print(args)

channel_sizes = [args.nhid] * args.levels
kernel_size = args.ksize
epoch_start = 0
model = O_TCN(input_channels, n_classes, channel_sizes, kernel_size=kernel_size, dropout=args.dropout, length=seq_length)
checkpoint = torch.load('weights/best.pt')
model.load_state_dict(checkpoint['model'])


def detect(opt):
    model.to(device).eval()
    with torch.no_grad():
        data = data_eval(seq_length=seq_length)
        if args.cuda:
            data = data.cuda()
        data = data.view(-1, input_channels, seq_length)
        output = model(data)
        y_pred = output.cpu().detach().numpy()
        # y_pred = sigmoid(y_pred)
        print('%s之后的5个工作日上涨概率为：%.2f%%\n' % (end_date, y_pred[0][0] * 100))


if __name__ == '__main__':
    torch.cuda.empty_cache()
    detect(args)
